MODULE fmm_qlm_builder

   USE fmm_global_consts
   USE fmm_global_types
   USE fmm_stats
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: fmm_get_raw_qlm,             &
             fmm_deallocate_qlm

   ! Public variables
   PUBLIC :: fmm_system_size, fmm_coord_shift

   ! Coordinate translation required to ensure all primitive (x,y,z) positive
   REAL(REALK), SAVE :: fmm_coord_shift(3)
   ! (minimum) dimension of a cube which encloses all moments
   REAL(REALK), SAVE :: fmm_system_size

   ! Number of moments, nuc charges, and/or potential grid points
   INTEGER(INTK), SAVE :: nelec, nnuc, ngrid
   ! Number of AO (contracted) basis functions
   INTEGER(INTK), SAVE :: nbas

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_get_raw_qlm(scheme,dens,LHS,RHS)

      IMPLICIT NONE
      TYPE(fmm_scheme),  INTENT(IN)  :: scheme
      REAL(REALK),       INTENT(IN)  :: dens(:,:)
      TYPE(raw_mm_data), INTENT(OUT) :: LHS, RHS
      TYPE(raw_mm_data) :: mm_elec, mm_nuc, mm_grid
      INTEGER(INTK) :: LMAX

      LMAX = scheme%raw_LMAX

      CALL fmm_get_dims_from_file(scheme,LMAX)

      ! Electronic part
      CALL fmm_allocate_mm_arrays(LMAX,nelec,mm_elec)
      CALL fmm_read_electronic_moms(scheme,dens,mm_elec)

      ! Nuclear part
!FIXME several problems if only allocate LMAX=0 further down route
      CALL fmm_allocate_mm_arrays(LMAX,nnuc,mm_nuc)
      CALL fmm_read_nuclear_moms(mm_nuc)

      ! Potential grid part
      CALL fmm_allocate_mm_arrays(0,ngrid,mm_grid)
      CALL fmm_read_potential_points(mm_grid)

      ! Assign LHS data
      SELECT CASE (scheme%LHS_mm_range)
         CASE (NUCLEAR_ONLY)
            CALL fmm_assign_LHS_RHS_data(LMAX,mm_nuc=mm_nuc,XHS=LHS)
         CASE (ELECTRONIC_ONLY)
            CALL fmm_assign_LHS_RHS_data(LMAX,mm_elec,XHS=LHS)
         CASE (ALL_MOMENTS)
            CALL fmm_assign_LHS_RHS_data(LMAX,mm_elec,mm_nuc,XHS=LHS)
         CASE (POTENTIAL_ONLY)
            CALL fmm_assign_LHS_RHS_data(0,mm_grid=mm_grid,XHS=LHS)
         CASE DEFAULT
            STOP 'fmm_distribute_LHS_RHS_data'
      END SELECT
      stat_raw_moms_LHS = SIZE(LHS%paras)
      print *, "stat_raw_moms_LHS =", SIZE(LHS%paras)

      ! Assign RHS data
      SELECT CASE (scheme%RHS_mm_range)
         CASE (NUCLEAR_ONLY)
            CALL fmm_assign_LHS_RHS_data(LMAX,mm_nuc=mm_nuc,XHS=RHS)
         CASE (ELECTRONIC_ONLY)
            CALL fmm_assign_LHS_RHS_data(LMAX,mm_elec,XHS=RHS)
         CASE (ALL_MOMENTS)
            CALL fmm_assign_LHS_RHS_data(LMAX,mm_elec,mm_nuc,XHS=RHS)
         CASE DEFAULT
            STOP 'fmm_distribute_LHS_RHS_data'
      END SELECT
      stat_raw_moms_RHS = SIZE(RHS%paras)
      print *, "stat_raw_moms_RHS =",  SIZE(RHS%paras)

      CALL fmm_deallocate_mm_arrays(mm_nuc)
      CALL fmm_deallocate_mm_arrays(mm_elec)
      CALL fmm_deallocate_mm_arrays(mm_grid)

   END SUBROUTINE fmm_get_raw_qlm

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_get_dims_from_file(scheme,LMAX_in)

      IMPLICIT NONE
      TYPE(fmm_scheme), INTENT(IN) :: scheme
      INTEGER(INTK),    INTENT(IN) :: LMAX_in

      INTEGER(INTK) :: LMAX, ntot

      ! Read number of electronic moments
      LMAX = LMAX_in
      nelec = 0
      IF ( (scheme%RHS_mm_range == ELECTRONIC_ONLY) .OR.   &
           (scheme%RHS_mm_range == ALL_MOMENTS)     .OR.   &
           (scheme%LHS_mm_range == ELECTRONIC_ONLY) .OR.   &
           (scheme%LHS_mm_range == ALL_MOMENTS) ) THEN
         OPEN (UNIT=IO2,FILE='multipoles_header.md1',STATUS='OLD',   &
               ACTION='READ',FORM='UNFORMATTED')
         REWIND (IO2)
         READ (IO2) LMAX, nbas, nelec
         CLOSE(UNIT=IO2,STATUS='KEEP')
         WRITE(IO1,'(A,I8)') "Number of electronic moments read = ", nelec
      END IF
      stat_n_basis = nbas

      ! Read number of nuclear charges
      nnuc = 0
      IF ( (scheme%RHS_mm_range == NUCLEAR_ONLY)  .OR.   &
           (scheme%RHS_mm_range == ALL_MOMENTS)   .OR.   &
           (scheme%LHS_mm_range == NUCLEAR_ONLY)  .OR.   &
           (scheme%LHS_mm_range == ALL_MOMENTS) ) THEN
         OPEN (UNIT=IO2,FILE='nuclei.fmm',STATUS='OLD',   &
               ACTION='READ',FORM='FORMATTED')
         REWIND (IO2)
         READ (IO2,'(I8)') nnuc
         CLOSE(UNIT=IO2,STATUS='KEEP')
         WRITE(IO1,'(A,I8)') "Number of nuclear charges read = ", nnuc
      END IF
print *, "nelec, nnuc", nelec, nnuc

      ! Read number of potential grid points
      ngrid = 0
      IF ( (scheme%LHS_mm_range == POTENTIAL_ONLY)   .OR.   &
           (scheme%RHS_mm_range == POTENTIAL_ONLY) ) THEN
         OPEN (UNIT=IO2,FILE='grid.fmm',STATUS='OLD',    &
               ACTION='READ',FORM='UNFORMATTED')
         REWIND (IO2)
         READ (IO2) ngrid
         CLOSE(UNIT=IO2,STATUS='KEEP')
         WRITE(IO1,'(A,I8)') "Number of potential grid points read = ", ngrid
      END IF

      ! Verifications
      ntot = nelec + nnuc + ngrid
      IF (ntot < 1) CALL fmm_quit('No moments generated!')
      IF (LMAX /= LMAX_in) THEN
         WRITE(IO1,*) LMAX, LMAX_in
         CALL fmm_quit('LMAX inconsistency in MM interface!')
      END IF

   END SUBROUTINE fmm_get_dims_from_file

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_allocate_mm_arrays(LMAX,ndim,mm_arr)

      IMPLICIT NONE
      INTEGER(INTK),     INTENT(IN)  :: LMAX, ndim
      TYPE(raw_mm_data), INTENT(OUT) :: mm_arr
      INTEGER(INTK) :: i, foo, alloc_error

      NULLIFY (mm_arr%paras, mm_arr%dens, mm_arr%J_indices, mm_arr%batch_map)
      NULLIFY (mm_arr%qlm, mm_arr%qlm_W, mm_arr%qlm_T)

      ALLOCATE (mm_arr%paras(ndim))
      ALLOCATE (mm_arr%J_indices(ndim))
      ALLOCATE (mm_arr%dens(ndim))

      ! Initialise
      DO i = 1, ndim 
         mm_arr%paras(i)%cntr = zero
         mm_arr%paras(i)%ext = zero
         mm_arr%paras(i)%id = 0
         mm_arr%paras(i)%batch = 0
         mm_arr%paras(i)%map_up = 0
         mm_arr%paras(i)%box = 0 
         mm_arr%paras(i)%bra = 0
         mm_arr%paras(i)%box_cntr = zero
         mm_arr%J_indices(i)%i_indx = 0
         mm_arr%J_indices(i)%j_indx = 0
         mm_arr%dens(i) = zero
      END DO

      ! Full set of Qlm moments only relevant for electronic part
      foo = (LMAX+1)**2
      WRITE(IO1,*) 'mm_arrays: Attempting to allocate',  &
                     MAX(1,foo*ndim*8/1000000), 'MB of memory...'
      ALLOCATE (mm_arr%qlm(foo,ndim), STAT=alloc_error)
      IF (alloc_error /= 0) WRITE(IO1,*) '... Failed!'

      mm_arr%qlm(:,:) = zero  ! only non-zero written explicitly

   END SUBROUTINE fmm_allocate_mm_arrays

!------------------------------------------------------------------------------

   SUBROUTINE fmm_deallocate_mm_arrays(mm_data)

      IMPLICIT NONE
      TYPE(raw_mm_data), INTENT(OUT) :: mm_data
      INTEGER(INTK) :: i

      IF (ASSOCIATED(mm_data%paras)) DEALLOCATE (mm_data%paras)
      IF (ASSOCIATED(mm_data%J_indices)) DEALLOCATE (mm_data%J_indices)
      IF (ASSOCIATED(mm_data%dens)) DEALLOCATE (mm_data%dens)
      IF (ASSOCIATED(mm_data%qlm)) DEALLOCATE (mm_data%qlm)
      IF (ASSOCIATED(mm_data%qlm_T)) DEALLOCATE (mm_data%qlm_T)
      IF (ASSOCIATED(mm_data%qlm_W)) DEALLOCATE (mm_data%qlm_W)

      IF (ASSOCIATED(mm_data%batch_map)) THEN
         DO i = 1, SIZE(mm_data%batch_map)
            CALL free_batch_map(mm_data%batch_map(i)%head)
         END DO
      END IF
      IF (ASSOCIATED(mm_data%batch_map)) DEALLOCATE(mm_data%batch_map)

      NULLIFY (mm_data%paras, mm_data%dens, mm_data%batch_map)
      NULLIFY (mm_data%qlm, mm_data%qlm_W, mm_data%qlm_T)
      NULLIFY (mm_data%J_indices)

   CONTAINS 

      RECURSIVE SUBROUTINE free_batch_map(node)

         IMPLICIT NONE
         TYPE(id_node), POINTER :: node
         IF (ASSOCIATED(node%next)) THEN
            CALL free_batch_map(node%next)
         END IF
         DEALLOCATE(node)
         NULLIFY(node)

      END SUBROUTINE free_batch_map

   END SUBROUTINE fmm_deallocate_mm_arrays

!-------------------------------------------------------------------------------
! Read in multipole moment data from interface file.
! In all this MM code we assume the order of moments is:
!   (0),(-1,0,1),(-2,-1,0,1,2)... wrt (L,M)

   SUBROUTINE fmm_read_electronic_moms(scheme,dens,mm_elec)

      IMPLICIT NONE
      TYPE(fmm_scheme),  INTENT(IN)  :: scheme
      REAL(REALK),       INTENT(IN)  :: dens(:,:)
      TYPE(raw_mm_data), INTENT(OUT) :: mm_elec
      REAL(REALK)   :: PX,PY,PZ, SPH
      INTEGER(INTK) :: I,J,L,M, A,B, LM, X

      ! Read electronic multipole moments into core
      IF (nelec == 0) RETURN
      OPEN (UNIT=IO2,FILE='multipoles.md1',STATUS='OLD',  &
            ACTION='READ',FORM='UNFORMATTED')
      REWIND (IO2)

      readloop: DO

         READ (IO2) I,L,M, A,B, PX,PY,PZ, SPH
         ! EOF marked by record with negative angular momentum
         IF (L < 0) EXIT readloop

         !IF ((L == 0) .AND. (ABS(SPH) > 1.0d-12) )   &
         IF (ABS(SPH) > 1.0d-12)  &
         WRITE (IO1,'(5I4,X,3F8.4,2E13.4)') I,L,M,A,B,PX,PY,PZ,SPH,dens(A,B)

         IF (A > nbas) CALL fmm_quit('interface file error 0')
         IF (B > nbas) CALL fmm_quit('interface file error 00')
         IF (I < 1) CALL fmm_quit('interface file error 1')
         IF (I > SIZE(mm_elec%qlm,2)) CALL fmm_quit('interface error 11')
         IF (((L+1)**2) > SIZE(mm_elec%qlm,1)) CALL fmm_quit('interface 156')

         ! Indices to map moments to orbitals in J-matrix
         mm_elec%J_indices(I)%i_indx = A
         mm_elec%J_indices(I)%j_indx = B

         ! Density matrix factor
         mm_elec%dens(I) = -dens(A,B)

         ! Multipole expansion centre
         mm_elec%paras(I)%cntr = (/ PX, PY, PZ /)

         ! See defn of 'extent' in p.424 MEST Helgaker et al
!FIXME NO EXTENT INFO 
         mm_elec%paras(I)%ext = 0

         ! Components (l,m) of multipole without density factorized in 
         LM = L*(L+1) +M +1 
         mm_elec%qlm(LM,I) = SPH

      END DO readloop

      CLOSE(UNIT=IO2,STATUS='KEEP')
      CALL fmm_get_system_size_and_shift(mm_elec%paras)
      IF (scheme%branch_free) CALL fmm_make_branch_free_extents(scheme,mm_elec)

   END SUBROUTINE fmm_read_electronic_moms

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_read_nuclear_moms(mm_nuc)

      IMPLICIT NONE
      TYPE(raw_mm_data), INTENT(OUT) :: mm_nuc
      REAL(REALK)   :: PX,PY,PZ
      INTEGER(INTK) :: Q, i

      ! Next read nuclei data: charge and location
      ! Positive charges give a positive potential contribution
      IF (nnuc == 0) RETURN
      OPEN (UNIT=IO2,FILE='nuclei.fmm',STATUS='OLD',  &
            ACTION='READ',FORM='FORMATTED')
      REWIND (IO2)
      READ (IO2,'(I8)') i

      WRITE (IO1,*) "Reading in point charges:"
      DO i = 1, nnuc 
         READ (IO2,'(F25.15)') PX
         READ (IO2,'(F25.15)') PY
         READ (IO2,'(F25.15)') PZ
         READ (IO2,'(I5)') Q
         WRITE (IO1,'(I4,X,3F8.4,I4)') I, PX,PY,PZ, Q
         mm_nuc%qlm(1,i) = DBLE(Q)
         mm_nuc%paras(i)%cntr = (/ PX, PY, PZ /)
      END DO

      mm_nuc%dens(:) = one

      CLOSE(UNIT=IO2,STATUS='KEEP')
      CALL fmm_get_system_size_and_shift(mm_nuc%paras)

   END SUBROUTINE fmm_read_nuclear_moms

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_read_potential_points(mm_grid)

      IMPLICIT NONE
      TYPE(raw_mm_data), INTENT(OUT) :: mm_grid
      REAL(REALK)   :: PX,PY,PZ, SPH
      INTEGER(INTK) :: i

      ! Read in the grid points when computing an arbitrary potential
      IF (ngrid == 0) RETURN
      OPEN (UNIT=IO2,FILE='grid.fmm',STATUS='OLD',  &
            ACTION='READ',FORM='UNFORMATTED')
      REWIND (IO2)
      READ (IO2) i

      DO i = 1, ngrid
         READ (IO2) PX,PY,PZ
!         WRITE (IO1,'(I4,X,3F8.4,2E13.4)') I, PX,PY,PZ
         mm_grid%paras(i)%cntr = (/ PX, PY, PZ /)
      END DO

      mm_grid%qlm(1,:) = one
      mm_grid%dens(:) = one

      CLOSE(UNIT=IO2,STATUS='KEEP')
      CALL fmm_get_system_size_and_shift(mm_grid%paras)

   END SUBROUTINE fmm_read_potential_points

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_get_system_size_and_shift(paras)

      IMPLICIT NONE
      TYPE(raw_mm_paras), INTENT(INOUT) :: paras(:)
      REAL(REALK)   :: sys_min(3), sys_max(3)
      INTEGER(INTK) :: i

      sys_min = paras(1)%cntr
      sys_max = paras(1)%cntr
      DO i = 1, SIZE(paras)
         sys_min(:) = MIN(sys_min(:),paras(i)%cntr(:))
         sys_max(:) = MAX(sys_max(:),paras(i)%cntr(:))
      END DO
      fmm_system_size = MAXVAL(sys_max - sys_min)
!      IF (fmm_system_size == zero) CALL fmm_quit('zero system size!')

      fmm_coord_shift = sys_min
!      DO i = 1, SIZE(paras)
!         paras(i)%cntr(:) = paras(i)%cntr(:) - sys_min(:)
!      END DO

   END SUBROUTINE fmm_get_system_size_and_shift

!-------------------------------------------------------------------------------
! Branch-free scheme basically uses CFMM code but we simplify all the
! extents such that they are all the same and only _very_ close interactions
! are avoided (for numerical stability)

   SUBROUTINE fmm_make_branch_free_extents(scheme,mm_data)

      IMPLICIT NONE
      TYPE(fmm_scheme),  INTENT(IN)  :: scheme
      TYPE(raw_mm_data), INTENT(OUT) :: mm_data

      INTEGER(INTK) :: i

      DO i = 1, SIZE(mm_data%paras)
         mm_data%paras(i)%ext = scheme%extent_min
      END DO

   END SUBROUTINE fmm_make_branch_free_extents

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_assign_LHS_RHS_data(LMAX,mm_elec,mm_nuc,mm_grid,XHS)

      IMPLICIT NONE
      INTEGER(INTK),               INTENT(IN)  :: LMAX
      TYPE(raw_mm_data), OPTIONAL, INTENT(IN)  :: mm_elec
      TYPE(raw_mm_data), OPTIONAL, INTENT(IN)  :: mm_nuc
      TYPE(raw_mm_data), OPTIONAL, INTENT(IN)  :: mm_grid
      TYPE(raw_mm_data),           INTENT(OUT) :: XHS

      INTEGER(INTK) :: ndim, lo,hi, i

      lo = 0
      hi = 0
      ndim = 0
      IF ( PRESENT(mm_nuc)  ) ndim = ndim + nnuc
      IF ( PRESENT(mm_elec) ) ndim = ndim + nelec
      IF ( PRESENT(mm_grid) ) ndim = ndim + ngrid
      IF (ndim == 0) CALL fmm_quit('fmm_assign_LHS_RHS_data')

      CALL fmm_allocate_mm_arrays(LMAX,ndim,XHS)

      IF ( PRESENT(mm_nuc) ) THEN
         hi = nnuc
         XHS%qlm(1,1:hi)     = mm_nuc%qlm(1,:)
         XHS%dens(1:hi)      = mm_nuc%dens(:)
         XHS%paras(1:hi)     = mm_nuc%paras(:)
         XHS%J_indices(1:hi) = mm_nuc%J_indices(:)
      END IF
      IF ( PRESENT(mm_elec) ) THEN
         lo = hi + 1
         hi = lo - 1 + nelec
         XHS%qlm(:,lo:hi)     = mm_elec%qlm(:,:)
         XHS%dens(lo:hi)      = mm_elec%dens(:)
         XHS%paras(lo:hi)     = mm_elec%paras(:)
         XHS%J_indices(lo:hi) = mm_elec%J_indices(:)
      END IF
      IF ( PRESENT(mm_grid) ) THEN
         lo = hi + 1
         hi = lo - 1 + ngrid
         XHS%qlm(1,lo:hi)     = mm_grid%qlm(1,:)
         XHS%dens(lo:hi)      = mm_grid%dens(:)
         XHS%paras(lo:hi)     = mm_grid%paras(:)
         XHS%J_indices(lo:hi) = mm_grid%J_indices(:)
      END IF

      NULLIFY(XHS%qlm_T)
      NULLIFY(XHS%qlm_W)

      ! Initialise parameter:moments mapping, 
      ! and batch numbers for unique centres
      DO i = 1, ndim
         XHS%paras(i)%id = i
         XHS%paras(i)%batch = i
      END DO

   END SUBROUTINE fmm_assign_LHS_RHS_data

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_deallocate_qlm(LHS,RHS)

      IMPLICIT NONE
      TYPE(raw_mm_data), INTENT(OUT) :: LHS, RHS

      CALL fmm_deallocate_mm_arrays(LHS)
      CALL fmm_deallocate_mm_arrays(RHS)

   END SUBROUTINE fmm_deallocate_qlm

!-------------------------------------------------------------------------------

END MODULE fmm_qlm_builder

